"""
内存记忆模块

基于内存的记忆存储实现
"""

import json
from datetime import datetime
from typing import Any, Dict, List, Optional, Union

from langchain_core.language_models import BaseLanguageModel
from langchain_core.messages import BaseMessage

from .base_memory import BaseMemory, MemoryItem, MemorySearchResult


class InMemoryStorage(BaseMemory):
    """
    内存记忆存储
    
    将记忆数据存储在内存中，适合临时使用和测试
    """
    
    def __init__(self, agent_id: str, llm: Optional[BaseLanguageModel] = None,
                 max_messages: int = 1000, auto_cleanup: bool = True):
        """
        初始化内存记忆存储
        
        Args:
            agent_id: 智能体ID
            llm: 大语言模型
            max_messages: 最大消息数量
            auto_cleanup: 是否自动清理旧消息
        """
        super().__init__(agent_id, llm)
        
        self.max_messages = max_messages
        self.auto_cleanup = auto_cleanup
        self._messages: Dict[str, MemoryItem] = {}
        self._message_order: List[str] = []  # 保持消息顺序
        self._search_index: Dict[str, List[str]] = {}  # 简单搜索索引
    
    def add_message(self, message: Union[str, BaseMessage], 
                   message_type: str = "human", 
                   metadata: Optional[Dict[str, Any]] = None) -> str:
        """
        添加消息到记忆
        
        Args:
            message: 消息内容
            message_type: 消息类型
            metadata: 元数据
            
        Returns:
            消息ID
        """
        # 标准化消息
        content = self._normalize_message(message)
        if isinstance(message, BaseMessage):
            message_type = self._get_message_type(message)
        
        # 生成消息ID
        message_id = self._generate_message_id()
        
        # 创建记忆项
        memory_item = MemoryItem(
            id=message_id,
            content=content,
            message_type=message_type,
            timestamp=datetime.now(),
            metadata=metadata or {},
            importance=0.5,
            tags=[]
        )
        
        # 存储消息
        self._messages[message_id] = memory_item
        self._message_order.append(message_id)
        
        # 更新搜索索引
        self._update_search_index(message_id, content)
        
        # 更新统计信息
        self._update_stats(content)
        
        # 自动清理
        if self.auto_cleanup and len(self._messages) > self.max_messages:
            self._cleanup_old_messages()
        
        return message_id
    
    def get_messages(self, limit: Optional[int] = None, 
                    message_type: Optional[str] = None) -> List[MemoryItem]:
        """
        获取消息列表
        
        Args:
            limit: 限制数量
            message_type: 消息类型过滤
            
        Returns:
            消息列表
        """
        # 按时间顺序获取消息
        message_ids = self._message_order.copy()
        
        # 应用类型过滤
        if message_type:
            message_ids = [
                msg_id for msg_id in message_ids
                if self._messages[msg_id].message_type == message_type
            ]
        
        # 应用数量限制
        if limit:
            message_ids = message_ids[-limit:]  # 获取最新的消息
        
        # 返回消息对象
        return [self._messages[msg_id] for msg_id in message_ids]
    
    def search_messages(self, query: str, limit: int = 10, 
                       threshold: float = 0.7) -> List[MemorySearchResult]:
        """
        搜索消息
        
        Args:
            query: 搜索查询
            limit: 限制数量
            threshold: 相关性阈值
            
        Returns:
            搜索结果列表
        """
        results = []
        query_lower = query.lower()
        query_words = query_lower.split()
        
        # 简单的关键词匹配
        for message_id, message in self._messages.items():
            content_lower = message.content.lower()
            
            # 计算匹配分数
            score = 0.0
            match_reasons = []
            
            # 完全匹配
            if query_lower in content_lower:
                score += 1.0
                match_reasons.append("完全匹配")
            
            # 单词匹配
            content_words = content_lower.split()
            matched_words = sum(1 for word in query_words if word in content_words)
            if matched_words > 0:
                word_score = matched_words / len(query_words)
                score += word_score * 0.8
                match_reasons.append(f"单词匹配({matched_words}/{len(query_words)})")
            
            # 标签匹配
            for tag in message.tags:
                if any(word in tag.lower() for word in query_words):
                    score += 0.3
                    match_reasons.append(f"标签匹配({tag})")
            
            # 元数据匹配
            for key, value in message.metadata.items():
                if isinstance(value, str) and any(word in value.lower() for word in query_words):
                    score += 0.2
                    match_reasons.append(f"元数据匹配({key})")
            
            # 检查是否达到阈值
            if score >= threshold:
                results.append(MemorySearchResult(
                    item=message,
                    relevance_score=score,
                    match_reason="; ".join(match_reasons)
                ))
        
        # 按相关性排序
        results.sort(key=lambda x: x.relevance_score, reverse=True)
        
        # 限制结果数量
        return results[:limit]
    
    def clear_memory(self) -> None:
        """清空记忆"""
        self._messages.clear()
        self._message_order.clear()
        self._search_index.clear()
        self._message_count = 0
        self._total_tokens = 0
        self.updated_at = datetime.now()
    
    def get_memory_stats(self) -> Dict[str, Any]:
        """
        获取记忆统计信息
        
        Returns:
            统计信息字典
        """
        # 统计消息类型
        type_counts = {}
        for message in self._messages.values():
            msg_type = message.message_type
            type_counts[msg_type] = type_counts.get(msg_type, 0) + 1
        
        # 统计标签
        tag_counts = {}
        for message in self._messages.values():
            for tag in message.tags:
                tag_counts[tag] = tag_counts.get(tag, 0) + 1
        
        return {
            "agent_id": self.agent_id,
            "total_messages": len(self._messages),
            "message_types": type_counts,
            "total_tokens": self._total_tokens,
            "memory_size": len(self._messages),
            "max_messages": self.max_messages,
            "auto_cleanup": self.auto_cleanup,
            "created_at": self.created_at.isoformat(),
            "updated_at": self.updated_at.isoformat(),
            "tag_counts": tag_counts,
            "search_index_size": len(self._search_index)
        }
    
    def update_message_importance(self, message_id: str, importance: float) -> bool:
        """
        更新消息重要性
        
        Args:
            message_id: 消息ID
            importance: 重要性评分 (0-1)
            
        Returns:
            是否更新成功
        """
        if message_id not in self._messages:
            return False
        
        self._messages[message_id].importance = max(0.0, min(1.0, importance))
        self.updated_at = datetime.now()
        return True
    
    def add_tags_to_message(self, message_id: str, tags: List[str]) -> bool:
        """
        为消息添加标签
        
        Args:
            message_id: 消息ID
            tags: 标签列表
            
        Returns:
            是否添加成功
        """
        if message_id not in self._messages:
            return False
        
        message = self._messages[message_id]
        for tag in tags:
            if tag not in message.tags:
                message.tags.append(tag)
        
        self.updated_at = datetime.now()
        return True
    
    def get_messages_by_tags(self, tags: List[str]) -> List[MemoryItem]:
        """
        根据标签获取消息
        
        Args:
            tags: 标签列表
            
        Returns:
            消息列表
        """
        results = []
        for message in self._messages.values():
            if any(tag in message.tags for tag in tags):
                results.append(message)
        
        # 按时间排序
        results.sort(key=lambda x: x.timestamp)
        return results
    
    def get_important_messages(self, min_importance: float = 0.7) -> List[MemoryItem]:
        """
        获取重要消息
        
        Args:
            min_importance: 最小重要性阈值
            
        Returns:
            重要消息列表
        """
        results = []
        for message in self._messages.values():
            if message.importance >= min_importance:
                results.append(message)
        
        # 按重要性排序
        results.sort(key=lambda x: x.importance, reverse=True)
        return results
    
    def export_memory(self, format: str = "json") -> Union[str, Dict]:
        """
        导出记忆数据
        
        Args:
            format: 导出格式 ("json" 或 "dict")
            
        Returns:
            导出的数据
        """
        data = {
            "agent_id": self.agent_id,
            "created_at": self.created_at.isoformat(),
            "updated_at": self.updated_at.isoformat(),
            "messages": [msg.to_dict() for msg in self._messages.values()],
            "message_order": self._message_order,
            "stats": self.get_memory_stats()
        }
        
        if format == "json":
            return json.dumps(data, ensure_ascii=False, indent=2)
        else:
            return data
    
    def import_memory(self, data: Union[str, Dict]) -> bool:
        """
        导入记忆数据
        
        Args:
            data: 导入的数据
            
        Returns:
            是否导入成功
        """
        try:
            if isinstance(data, str):
                data = json.loads(data)
            
            # 清空现有记忆
            self.clear_memory()
            
            # 导入消息
            for msg_data in data.get("messages", []):
                message = MemoryItem.from_dict(msg_data)
                self._messages[message.id] = message
            
            # 恢复消息顺序
            self._message_order = data.get("message_order", [])
            
            # 重建搜索索引
            self._rebuild_search_index()
            
            # 更新统计信息
            self._message_count = len(self._messages)
            self._total_tokens = sum(len(msg.content.split()) for msg in self._messages.values())
            
            return True
            
        except Exception as e:
            print(f"导入记忆数据失败: {e}")
            return False
    
    def _update_search_index(self, message_id: str, content: str) -> None:
        """更新搜索索引"""
        words = content.lower().split()
        for word in words:
            if len(word) > 2:  # 只索引长度大于2的词
                if word not in self._search_index:
                    self._search_index[word] = []
                if message_id not in self._search_index[word]:
                    self._search_index[word].append(message_id)
    
    def _rebuild_search_index(self) -> None:
        """重建搜索索引"""
        self._search_index.clear()
        for message_id, message in self._messages.items():
            self._update_search_index(message_id, message.content)
    
    def _cleanup_old_messages(self) -> None:
        """清理旧消息"""
        if len(self._messages) <= self.max_messages:
            return
        
        # 按重要性排序，保留重要的消息
        messages_by_importance = sorted(
            self._messages.items(),
            key=lambda x: (x[1].importance, x[1].timestamp),
            reverse=True
        )
        
        # 保留最重要的消息
        keep_count = int(self.max_messages * 0.8)  # 保留80%
        messages_to_keep = messages_by_importance[:keep_count]
        
        # 清空并重建
        self._messages.clear()
        self._message_order.clear()
        
        for message_id, message in messages_to_keep:
            self._messages[message_id] = message
            self._message_order.append(message_id)
        
        # 重建搜索索引
        self._rebuild_search_index()
    
    def get_messages_by_time_range(self, start_time: datetime, 
                                  end_time: Optional[datetime] = None) -> List[MemoryItem]:
        """
        根据时间范围获取消息
        
        Args:
            start_time: 开始时间
            end_time: 结束时间
            
        Returns:
            消息列表
        """
        if end_time is None:
            end_time = datetime.now()
        
        results = []
        for message in self._messages.values():
            if start_time <= message.timestamp <= end_time:
                results.append(message)
        
        # 按时间排序
        results.sort(key=lambda x: x.timestamp)
        return results
    
    def get_message_by_id(self, message_id: str) -> Optional[MemoryItem]:
        """
        根据ID获取消息
        
        Args:
            message_id: 消息ID
            
        Returns:
            消息对象或None
        """
        return self._messages.get(message_id)
    
    def delete_message(self, message_id: str) -> bool:
        """
        删除消息
        
        Args:
            message_id: 消息ID
            
        Returns:
            是否删除成功
        """
        if message_id not in self._messages:
            return False
        
        # 删除消息
        del self._messages[message_id]
        
        # 从顺序列表中移除
        if message_id in self._message_order:
            self._message_order.remove(message_id)
        
        # 从搜索索引中移除
        for word, message_ids in self._search_index.items():
            if message_id in message_ids:
                message_ids.remove(message_id)
                if not message_ids:
                    del self._search_index[word]
        
        self.updated_at = datetime.now()
        return True
    
    def get_conversation_summary(self, max_messages: int = 10) -> str:
        """
        获取对话摘要
        
        Args:
            max_messages: 最大消息数量
            
        Returns:
            对话摘要
        """
        recent_messages = self.get_messages(limit=max_messages)
        
        if not recent_messages:
            return "没有对话历史"
        
        summary_parts = []
        for msg in recent_messages:
            role = "用户" if msg.message_type == "human" else "AI"
            summary_parts.append(f"{role}: {msg.content[:100]}...")
        
        return "\n".join(summary_parts)
    
    def __len__(self) -> int:
        """返回消息数量"""
        return len(self._messages)
    
    def __contains__(self, message_id: str) -> bool:
        """检查是否包含指定消息"""
        return message_id in self._messages
    
    def __iter__(self):
        """迭代所有消息"""
        return iter(self._messages.values())


# 注册内存记忆类型
from .base_memory import _memory_manager
_memory_manager.register_memory_type("in_memory", InMemoryStorage)
_memory_manager.register_memory_type("buffer", InMemoryStorage)  # 别名
